#!/usr/bin/python3
import os
import cv2
import numpy as np
import pandas as pd

face_attr_list_file = r'/rootfs/media/kasim/DataSet/FaceVideo/CropImageMargin/file_bbox_list_0.5.attr.txt'

# VERY_GOOD = 0
# GOOD = 1
# BAD = 2
# VERY_BAD = 3
filter_quality = 3
filter_quality_thr = 0.25

min_pitch = -45
max_pitch = 45
min_roll = -45
max_roll = 45
min_yaw = -45
max_yaw = 45

min_angle_score = 0.1
max_angle_diff = 20


def main():
    # file_name;bin_pitch;bin_roll;bin_yaw;pitch;roll;yaw;quality
    attr_cvs = pd.read_csv(face_attr_list_file, sep=';')
    file_names = attr_cvs['file_name']
    qualitys = attr_cvs['quality']
    bin_pitchs = attr_cvs['bin_pitch']
    bin_rolls = attr_cvs['bin_roll']
    bin_yaws = attr_cvs['bin_yaw']
    pitchs = attr_cvs['pitch']
    rolls = attr_cvs['roll']
    yaws = attr_cvs['yaw']

    file_count = 0
    dir_file_selects = {}
    for file_name, quality_s, bin_pitch_s, bin_roll_s, bin_yaw_s, pitch_s, roll_s, yaw_s in zip(file_names, qualitys, bin_pitchs, bin_rolls, bin_yaws, pitchs, rolls, yaws):
        file_count += 1
        if file_count % 10000 == 0:
            print('Process Filter File Count: {}'.format(file_count))
        quality = list(map(float, quality_s.split(',')))
        quality = np.array(quality, dtype=np.float32)
        max_idx = np.argmax(quality)
        if (max_idx == filter_quality) and (quality[max_idx] >= filter_quality_thr):
            continue

        pitch = float(pitch_s)
        roll = float(roll_s)
        yaw = float(yaw_s)

        bin_pitch = bin_pitch_s.split(',')
        bin_roll = bin_roll_s.split(',')
        bin_yaw = bin_yaw_s.split(',')

        bin_pitch_score = float(bin_pitch[1])
        bin_pitch = float(bin_pitch[0])
        bin_roll_score = float(bin_roll[1])
        bin_roll = float(bin_roll[0])
        bin_yaw_score = float(bin_yaw[1])
        bin_yaw = float(bin_yaw[0])

        pitch_diff = abs(bin_pitch - pitch)
        roll_diff = abs(bin_roll - roll)
        yaw_diff = abs(bin_yaw - yaw)

        is_select1 = False
        is_select2 = False
        is_select3 = False
        if pitch < min_pitch or max_pitch < pitch or roll < min_roll or max_roll < roll or yaw < min_yaw or max_yaw < yaw:
            is_select1 = True

        if max_angle_diff <= pitch_diff or max_angle_diff <= roll_diff or max_angle_diff <= yaw_diff:
            is_select2 = True

        if bin_pitch_score <= min_angle_score or bin_roll_score <= min_angle_score or bin_yaw_score <= min_angle_score:
            is_select3 = True

        dir_name = os.path.dirname(file_name)
        base_name = os.path.basename(file_name)
        file_selects = dir_file_selects.get(dir_name, None)
        if file_selects is None:
            file_selects = [[], []]
            dir_file_selects[dir_name] = file_selects

        file_selects[0].append(base_name)
        # # if is_select1 or is_select2 or is_select3:
        # if is_select1 or (is_select2 and is_select3):
        # # if is_select1 and is_select2 and is_select3:
        #     # select file
        #     file_selects[1].append('1')
        # else:
        #     # no select file
        #     file_selects[1].append('0')
        file_selects[1].append('0')
    print('Process Filter File Count: {}'.format(file_count))

    file_count = 0
    out_file_path = os.path.join(os.path.dirname(face_attr_list_file), 'select_images.txt')
    with open(out_file_path, 'w') as out_file:
        for dir_name, file_selects in dir_file_selects.items():
            file_count += 1
            if file_count % 1000 == 0:
                print('Process Select File Count: {}'.format(file_count))
            files = file_selects[0]
            selects = file_selects[1]
            files = ' '.join(files)
            selects = ' '.join(selects)
            out_info_str = '{} {};{}\n'.format(dir_name, files, selects)
            out_file.write(out_info_str)
    os.system('chmod a+wr {}'.format(out_file_path))
    print('Process Select File Count: {}'.format(file_count))


if __name__ == '__main__':
    main()
